from typing import List
from omegaconf import DictConfig

import torch
import numpy as np

from torch.utils.data import Dataset

from ecgcmr.signal.sig_augmentations.ECGAugmentations import ECGAugmentations
from ecgcmr.imaging.img_augmentations.ImageAugmentations import ImageAugmentations


class MultiModalImagingECGDataset(Dataset):
    """
    Multimodal dataset that generates views of imaging and ECG data for contrastive learning.
    """
    def __init__(self, cfg: DictConfig, mode: str, use_peaks_location: bool = False, apply_augmentations: bool = True) -> None:
        self.cfg = cfg
        self.mode = mode
        self.apply_augmentations = apply_augmentations
        self.use_peaks_location = use_peaks_location

        self._load_data()
        self._initialize_augmentations()

    def _load_data(self):
        paths = self.cfg.dataset.paths
        data_mapping = {
            'train': (paths.data_imaging_train, paths.data_ecg_train, paths.data_patient_ids_train),
            'val': (paths.data_imaging_val, paths.data_ecg_val, paths.data_patient_ids_val)
        }

        if self.mode in data_mapping:
            self.data_imaging, self.data_ecg, self.data_patients_ids = [
                np.load(path, mmap_mode='c') for path in data_mapping[self.mode]
            ]
            if self.use_peaks_location:
                self.data_ecg_r_peaks = np.load(getattr(paths, f'data_ecg_rpeaks_updated_{self.mode}'), allow_pickle=True)
                self.data_imaging_tpf = np.load(getattr(paths, f'data_imaging_time_per_frame_{self.mode}'))

    def _initialize_augmentations(self):
        self.image_augmentations = ImageAugmentations(
            cfg=self.cfg.augmentations.imaging,
            img_size=self.cfg.dataset.img_size,
            use_peaks_location=self.use_peaks_location,
            apply_augmentations=self.apply_augmentations
        )
        self.ecg_augmentations = ECGAugmentations(
            cfg=self.cfg.augmentations.ecg,
            sampling_rate=self.cfg.dataset.sampling_rate,
            input_electrodes=self.cfg.dataset.input_electrodes,
            single_modality=self.cfg.dataset.single_modality,
            use_peaks_location=self.use_peaks_location,
            seed=self.cfg.seed,  
            apply_augmentations=self.apply_augmentations
        )

    def generate_ecg_views(self, index: int) -> List[torch.Tensor]:
        """
        Generates views of ECG data.
        """
        ecg_orig = self.data_ecg[index]

        if self.use_peaks_location:
            ecg_data_dict = {"ecg": ecg_orig}
            ecg_data_dict.update({"rpeaks": self.data_ecg_r_peaks[index]})
        else:
            ecg_data_dict = ecg_orig

        ecg_data = self.ecg_augmentations(ecg_data_dict)

        if self.use_peaks_location:
            return {"aug_ecg": ecg_data["ecg"], 
                    "rpeaks": ecg_data["rpeaks"]}
        else:
            return {"aug_ecg": ecg_data}

    def generate_imaging_views(self, index: int) -> List[torch.Tensor]:
        """
        Generates views of a subject's imaging data.
        """
        img_orig = self.data_imaging[index]
        img_data = {"image": img_orig}

        img_aug = self.image_augmentations(img_data)
        
        img_data_dict = {
            "img_aug": img_aug,
        }

        if self.use_peaks_location:
            img_data_dict.update({"tpf": self.data_imaging_tpf[index]})
        
        return img_data_dict

    def __getitem__(self, index: int):
        ecg_info = self.generate_ecg_views(index)
        img_info = self.generate_imaging_views(index)
        patient_id = self.data_patients_ids[index]

        result = {
            "ecg_aug": ecg_info["aug_ecg"],
            "image_aug": img_info["img_aug"],
            "patient_id": patient_id
        }

        if self.use_peaks_location:
            result.update({"rpeaks": ecg_info["rpeaks"]})
            result.update({"tpf": img_info["tpf"]})

        return result

    def __len__(self) -> int:
        return len(self.data_ecg)